import numpy as np
import pandas as pd
from sklearn.preprocessing import MinMaxScaler
from torch.utils.data import Dataset, DataLoader
import torch


class StockDataset(Dataset):
    def __init__(self, X, y):
        self.X = torch.FloatTensor(X)
        self.y = torch.FloatTensor(y)

    def __len__(self):
        return len(self.X)

    def __getitem__(self, idx):
        return self.X[idx], self.y[idx]


def create_sequences(data, lookback):
    """创建时间序列数据"""
    X, y = [], []
    for i in range(len(data) - lookback):
        X.append(data[i:(i + lookback)])
        y.append(data[i + lookback])
    return np.array(X), np.array(y)


def prepare_stock_data(file_path, lookback=60, train_ratio=0.8):
    """
    准备股票数据
    Args:
        file_path: CSV文件路径
        lookback: 回看天数
        train_ratio: 训练集比例
    """
    # 读取数据
    df = pd.read_csv(file_path)
    # 使用收盘价
    data = df['Close'].values.reshape(-1, 1)

    # 数据标准化
    scaler = MinMaxScaler()
    data_scaled = scaler.fit_transform(data)

    # 创建序列
    X, y = create_sequences(data_scaled, lookback)

    # 划分训练集和测试集
    train_size = int(len(X) * train_ratio)
    X_train, X_test = X[:train_size], X[train_size:]
    y_train, y_test = y[:train_size], y[train_size:]

    # 创建数据集
    train_dataset = StockDataset(X_train, y_train)
    test_dataset = StockDataset(X_test, y_test)

    return train_dataset, test_dataset, scaler
